import torch
import torch.nn as nn

# Generator
innerChannels = 32


class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d(in_channels=100,
                               out_channels=innerChannels * 16,
                               kernel_size=(12, 4)),
            nn.BatchNorm2d(innerChannels * 16),
            nn.ReLU(True),
            nn.ConvTranspose2d(in_channels=innerChannels * 16,
                               out_channels=innerChannels * 8,
                               kernel_size=(4, 4),
                               stride=(2, 2),
                               padding=(1, 1)),
            nn.BatchNorm2d(innerChannels * 8),
            nn.ReLU(True),
            nn.ConvTranspose2d(in_channels=innerChannels * 8,
                               out_channels=innerChannels * 4,
                               kernel_size=(4, 4),
                               stride=(2, 2),
                               padding=(1, 1)),
            nn.BatchNorm2d(innerChannels * 4),
            nn.ReLU(True),
            nn.ConvTranspose2d(in_channels=innerChannels * 4,
                               out_channels=innerChannels * 2,
                               kernel_size=(4, 4),
                               stride=(2, 2),
                               padding=(1, 1)),
            nn.BatchNorm2d(innerChannels * 2),
            nn.ReLU(True),
            nn.ConvTranspose2d(in_channels=innerChannels * 2,
                               out_channels=innerChannels,
                               kernel_size=(4, 4),
                               stride=(2, 2),
                               padding=(1, 1)),
            nn.BatchNorm2d(innerChannels),
            nn.ReLU(True),
            nn.ConvTranspose2d(in_channels=innerChannels,
                               out_channels=3,
                               kernel_size=(4, 4),
                               stride=(2, 2),
                               padding=(1, 1)),
            nn.Tanh())

    def forward(self, inputs):
        return self.main(inputs)


if __name__ == '__main__':
    model = Generator().cuda()
    x = torch.ones((1, 100, 1, 1), dtype=torch.float32).cuda()
    y = model(x)
    print(y.shape)
